﻿' This software is licensed under Microsoft Public License (Ms-PL), which governs use of the accompanying software. 
' If you use the software, you accept this license. 

' The software is licensed "as-is." You bear the risk of using it. The contributors give no express warranties, 
' guarantees or conditions. You may have additional consumer rights under your local laws which this license cannot 
' change. To the extent permitted under your local laws, the contributors exclude the implied warranties of 
' merchantability, fitness for a particular purpose and non-infringement.

Imports System.ServiceModel.Web
Imports System.ServiceModel.Rest
Imports System.ServiceModel.Channels
Imports System.ServiceModel.Dispatcher

Namespace Dispatcher
    ''' <summary>
    ''' Inspects REST messages for cookie management.
    ''' </summary>
    ''' <remarks></remarks>
    Public Class RestCookieInspector
        Implements IClientMessageInspector

        Private _cookie As String

        ''' <summary>
        ''' Enables inspection or modification of a message after a reply message is received.
        ''' </summary>
        ''' <param name="reply">Reply recieved from the service.</param>
        ''' <param name="correlationState"></param>
        ''' <remarks></remarks>
        Public Sub AfterReceiveReply(ByRef reply As System.ServiceModel.Channels.Message, correlationState As Object) Implements System.ServiceModel.Dispatcher.IClientMessageInspector.AfterReceiveReply

            If reply.Properties.ContainsKey(HttpResponseMessageProperty.Name) Then

                Dim response As HttpResponseMessageProperty = reply.Properties(HttpResponseMessageProperty.Name)
                If response IsNot Nothing Then

                    Dim cookie As String = response.Headers.Get("Set-Cookie")
                    If Not String.IsNullOrEmpty(cookie) Then
                        _cookie = cookie
                    End If
                End If
            End If
        End Sub

        ''' <summary>
        ''' Enables inspection or modification of a message before a request message is send. 
        ''' </summary>
        ''' <param name="request">Request to send.</param>
        ''' <param name="channel">Channel to send the request through.</param>
        ''' <returns></returns>
        ''' <remarks></remarks>
        Public Function BeforeSendRequest(ByRef request As System.ServiceModel.Channels.Message, channel As System.ServiceModel.IClientChannel) As Object Implements System.ServiceModel.Dispatcher.IClientMessageInspector.BeforeSendRequest

            If (Not String.IsNullOrEmpty(_cookie) AndAlso request.Properties.ContainsKey("httpRequest")) Then
                Dim httpRequest As HttpRequestMessageProperty = request.Properties("httpRequest")

                If httpRequest IsNot Nothing Then
                    httpRequest.Headers("Cookie") = _cookie
                End If
            End If

            Return Nothing
        End Function
    End Class
End Namespace
